#include <cmath>
#include <dirent.h>
#include <string.h>
#include <map>
#include "acllite_dvpp_lite/ImageProc.h"
#include "acllite_om_execute/ModelProc.h"
#include "opencv2/opencv.hpp"

using namespace std;
using namespace acllite;
using namespace cv;

typedef struct BoundBox {
	float x;
	float y;
	float width;
	float height;
	float score;
	int classIndex;
} BoundBox;
float iou(BoundBox box1, BoundBox box2)
{
	float xLeft = max(box1.x, box2.x);
	float yTop = max(box1.y, box2.y);
	float xRight = min(box1.x + box1.width, box2.x + box2.width);
	float yBottom = min(box1.y + box1.height, box2.y + box2.height);
	float width = max(0.0f, xRight - xLeft);
	float hight = max(0.0f, yBottom - yTop);
	float area = width * hight;
	float iou =  area / (box1.width * box1.height + box2.width * box2.height - area);
	return iou;
}
bool sortScore(BoundBox box1, BoundBox box2)
{
	return box1.score > box2.score;
}
int main()
{    
    vector<string> labels = { {"BACKGROUND"},{"with_mask"},{"mask_weared_incorrect"},{"without_mask"}};
    AclLiteResource aclResource;
    bool ret = aclResource.Init();
    CHECK_RET(ret, LOG_PRINT("[ERROR] InitACLResource failed."); return 1);
    
    ImageProc imageProc;
    ModelProc modelProc;
    ret = modelProc.Load("../model/mobilenet-ssd.om");
    CHECK_RET(ret, LOG_PRINT("[ERROR] load model mobilenet-ssd.om failed."); return 1);
	string imagePath = "../data/mask.jpg";
    ImageData src = imageProc.Read(imagePath);

    CHECK_RET(src.size, LOG_PRINT("[ERROR] ImRead image failed."); return 1);
    ImageData dst;
    ImageSize dsize(304, 300);

    imageProc.Resize(src, dst, dsize);    
    ret = modelProc.CreateInput(static_cast<void *>(dst.data.get()), dst.size);
    CHECK_RET(ret, LOG_PRINT("[ERROR] Create model input failed."); return 1);
    vector<InferenceOutput> inferOutputs;
    ret = modelProc.Execute(inferOutputs);
    CHECK_RET(ret, LOG_PRINT("[ERROR] model execute failed."); return 1);
    uint32_t dataSize = inferOutputs[0].size;
	uint32_t size = inferOutputs[1].size;
    // get result from output data set
    float* scores = static_cast<float*>(inferOutputs[0].data.get());
	float* boxes = static_cast<float*>(inferOutputs[1].data.get());
    if (scores == nullptr || boxes == nullptr) {
        LOG_PRINT("get result from output data set failed.");
        return 1;
    }
	size_t classNum = 4; 
    size_t boxes_nums = 3000;
	size_t candidate_size = 200;
	size_t top_k = 20;
	float prob_threshold = 0.7;
	float iou_threshold = 0.45;
    const double fountScale = 0.5;
    const uint32_t lineSolid = 2;
    const uint32_t labelOffset = 11;
    const cv::Scalar fountColor(0, 0, 255);
    const vector <cv::Scalar> colors{
        cv::Scalar(237, 149, 100), cv::Scalar(0, 215, 255),
        cv::Scalar(50, 205, 50), cv::Scalar(139, 85, 26)};
	cv::Mat srcImage = cv::imread(imagePath);
	int width = srcImage.cols;
    int height = srcImage.rows;
    for(int index = 1; index < classNum; index++) {
		vector<BoundBox> box_scores;
		vector<BoundBox> result;
		for(int j = 0; j < boxes_nums; ++j){
			if(scores[j * classNum + index] > prob_threshold){	
				BoundBox box;
				box.score = scores[j * classNum + index];
				box.width = (boxes[4 * j + 2]- boxes[4 * j]) * width ;
				box.height = (boxes[4 * j + 3] - boxes[4 * j + 1]) * height;
				box.x = boxes[4 * j] * width;
				box.y = boxes[4 * j + 1] * height;
				box.classIndex = index;
				box_scores.push_back(box);
			}
		}
		std::sort(box_scores.begin(),box_scores.end(),sortScore);
		if(box_scores.size() > candidate_size){
			box_scores.erase(box_scores.begin() + candidate_size + 1, box_scores.end());
		}
		int len = box_scores.size();
		if(len > 0){
			for(int i = 0;i < box_scores.size(); i++){
				if(result.size() == top_k) break;
				result.push_back(box_scores[i]);
				for(int j = i + 1; j < box_scores.size();j++){
					float iou_t = iou(box_scores[i],box_scores[j]);
					if(iou_t > iou_threshold){
						box_scores.erase(box_scores.begin() + j);
						j--;
					}
				}
			}	
		}
		for (size_t i = 0; i < result.size(); ++i) {
			cv::Point leftUpPoint, rightBottomPoint;
			leftUpPoint.x = result[i].x ;
			leftUpPoint.y = result[i].y;
			rightBottomPoint.x = result[i].x + result[i].width;
			rightBottomPoint.y = result[i].y + result[i].height;
			cv::rectangle(srcImage, leftUpPoint, rightBottomPoint, colors[i % colors.size()], lineSolid);
			string className = labels[result[i].classIndex];
			string markString = to_string(result[i].score) + ":" + className;
			cv::putText(srcImage, markString, cv::Point(leftUpPoint.x, leftUpPoint.y + labelOffset),
			cv::FONT_HERSHEY_COMPLEX, fountScale, fountColor);
		}
		
	}
	string savePath = "../output/out_0.jpg";
	cv::imwrite(savePath, srcImage);
    return 0;
}

